import torch
from torch import nn
from d2l import torch as d2l
import Attention as at

def transpose_qkv(X,num_heads):
    X = X.reshape(X.shape[0],X.shape[1],num_heads,-1)
    X = X.permute(0,2,1,3)
    return X.reshape(-1,X.shape[2],X.shape[3])

def transpose_output(X,num_heads):
    X = X.reshape(-1,num_heads,X.shape[1],X.shape[2])
    X = X.permute(0,2,1,3)
    return X.reshape(X.shape[0],X.shape[1],-1)

class MultiHeadAttention(nn.Module):
    def __init__(self, key_size,query_size,value_size,
                 num_hiddens,num_heads,dropout,bias=True, **kwargs):
        super().__init__( **kwargs)
        self.num_heads = num_heads
        self.attention = at.DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size,num_hiddens,bias=bias)
        self.W_k = nn.Linear(key_size,num_hiddens,bias=bias)
        self.W_v = nn.Linear(value_size,num_hiddens,bias=bias)
        self.W_o = nn.Linear(num_hiddens,num_hiddens,bias=bias)
    
    def forward(self,queries,keys,values,valid_lens):
        queries = transpose_qkv(self.W_q(queries),self.num_heads)
        keys = transpose_qkv(self.W_k(keys),self.num_heads)
        values = transpose_qkv(self.W_v(values),self.num_heads)
        if valid_lens is not None:
            valid_lens = torch.repeat_interleave(valid_lens,repeats=self.num_heads,dim=0)
        output = self.attention(queries,keys,values,valid_lens)
        output_concat = transpose_output(output,self.num_heads)
        return self.W_o(output_concat)

import math
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l

class PositionWiseFFN(nn.Module):
    def __init__(self, num_input,num_hiddens,num_output, **kwargs):
        super().__init__(**kwargs)
        self.dense1 = nn.Linear(num_input,num_hiddens)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(num_hiddens,num_output)
    
    def forward(self,X):
        return self.dense2(self.relu(self.dense1(X)))

class AddNorm(nn.Module):
    def __init__(self,normalized_shape,dropout,**kwargs):
        super().__init__(**kwargs)
        self.dropout = dropout
        self.ln = nn.LayerNorm(normalized_shape)
    
    def forward(self,X,Y):
        return self.ln(self.dropout(Y) + X)

class EncoderBlock(nn.Module):
    def __init__(self,key_size,query_size,value_size,
                 num_hiddens,norm_shape,ffn_num_input,ffn_num_hiddens,num_heads,
                 dropout,use_bias=False,**kwargs):
        super().__init__(**kwargs)
        self.attention = MultiHeadAttention(key_size,query_size,value_size,
                                            num_hiddens,num_heads,dropout,use_bias)
        self.addnorm1 = AddNorm(norm_shape,dropout)
        self.ffn = PositionWiseFFN(ffn_num_input,ffn_num_hiddens,num_hiddens)
        self.addnorm2 = AddNorm(norm_shape,dropout)
    
    def forward(self,X,valid_lens):
        Y = self.addnorm1(X,self.attention(X,X,X,valid_lens))
        return self.addnorm2(Y,self.ffn(Y))

import SelfAttention as sf
class TransformerEncoder(d2l.Encoder):
    def __init__(self,vocab_size,
                 key_size,query_size,value_size,
                 num_hiddens,norm_shape,
                 ffn_num_input,ffn_num_hiddens,
                 num_heads,num_layers,dropout,use_bias=False,**kwargs):
        super().__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.embedding = nn.Embedding(vocab_size,num_hiddens)
        self.pos_encoding = sf.PositionalEncoding(num_hiddens,dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module(
                "block" + str(i),
                EncoderBlock(key_size,query_size,value_size,
                             num_hiddens,norm_shape,
                             ffn_num_input,ffn_num_hiddens,
                             num_heads,dropout,use_bias)
            )
    
    def forward(self,X,valid_lens,*args):
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self.attention_weights = [None] * len(self.blks)
        for i,blk in enumerate(self.blks):
            X = blk(X,valid_lens)
            self.attention_weights[i] = blk.attention.attention_weights
        return X
    
class DecoderBlock(nn.Module):
    def __init__(self,key_size,query_size,value_size,
                 num_hiddens,norm_shape,
                 ffn_num_input,ffn_num_hiddens,num_heads,
                 dropout,i,**kwargs):
        super().__init__(**kwargs)
        self.i = i
        self.attention1 = MultiHeadAttention(key_size,query_size,value_size,
                                             num_hiddens,num_heads,dropout)
        self.addnorm1 = AddNorm(norm_shape,dropout)
        self.attention2 = MultiHeadAttention(key_size,query_size,value_size,
                                             num_hiddens,num_heads,dropout)
        self.addnorm2 = AddNorm(norm_shape,dropout)
        self.ffn = PositionWiseFFN(ffn_num_input,ffn_num_hiddens,num_hiddens)
        self.addnorm3 = AddNorm(norm_shape,dropout)
    
    def forward(self,X,state):
        enc_outputs,enc_valid_lens = state[0],state[1]
        if state[2][self.i] is not None:
            key_values = X
        else:
            key_values = torch.cat((state[2][self.i],X),axis=1)
        state[2][self.i] = key_values
        if self.training:
            batch_size,num_steps,_ = X.shape
            dec_valid_lens = torch.arange(1,num_steps + 1,device=X.device).repeat(batch_size,1)
        else:
            dec_valid_lens = None
        X2 = self.attention1(X,key_values,key_values,dec_valid_lens)
        Y = self.addnorm1(X,X2)
        Y2 = self.attention2(Y,enc_outputs,enc_outputs,enc_valid_lens)
        Z = self.addnorm2(Y,Y2)
        return self.addnorm3(X,self.ffn(Z)),state

class TransformerDecoder(d2l.AttentionDecoder):
    def __init__(self,vocab_size,
                 key_size,query_size,value_size,
                 num_hiddens,norm_shape,
                 ffn_num_input,ffn_num_hiddens,
                 num_heads,num_layers,dropout,**kwargs):
        super().__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size,num_hiddens)
        self.pos_encoding = sf.PositionalEncoding(num_hiddens,dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module(
                "block" + str(i),
                DecoderBlock(key_size,query_size,value_size,
                             num_hiddens,norm_shape,
                             ffn_num_input,ffn_num_hiddens,
                             num_heads,dropout,i)
            )
        self.dense = nn.Linear(num_hiddens,vocab_size)
    
    def init_state(self,enc_outputs,enc_valid_lens):
        return [enc_outputs,enc_valid_lens,[None] * self.num_layers]
    
    def forward(self,X,state):
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self._attention_weights = [[None] * len(self.blks) for _ in range(2)]
        for i,blk in enumerate(self.blks):
            X,state = blk(X,state)
            self._attention_weights[0][i] = blk.attention1.attention_weights
            self._attention_weights[1][i] = blk.attention2.attention_weights
        return self.dense(X),state
    
    @property
    def attention_weights(self):
        return self._attention_weights
    
